import cv2
import numpy as np
import math
import os
import csv
import random
from tqdm import tqdm


class GenDataset:
    def __init__(
        self,
        scenes_dir: str = "dataset_tools/sources/mjcf_scenes",
        dataset_save_dir: str = "dataset",
        resolution: float = 0.02,
        dataset_pre_scene: int = 100,
        scene_size: list = [20, 20],
        roi_imge_width: int = 250
    ):
        self.scenes_dir = scenes_dir
        self.dataset_save_dir = dataset_save_dir
        self.resolution = resolution
        self.dataset_pre_scene = dataset_pre_scene
        self.scene_size = scene_size
        self.scene_size_pixel = [
            self.scene_size[0] / self.resolution,
            self.scene_size[1] / self.resolution
        ]
        self.roi_imge_width = roi_imge_width
        self.scan_width = self.roi_imge_width * self.resolution

    def extract_rotated_square(
        self, 
        image_path, 
        center_x, center_y, 
        side_length, 
        angle_deg, 
        output_path
    ):
        """
        从图像中截取一个旋转的正方形区域
        参数:
            image_path: 输入图像路径
            center_x: 正方形中心x坐标
            center_y: 正方形中心y坐标
            side_length: 正方形边长
            angle_deg: 旋转角度(度)
            output_path: 输出图像路径
        """
        # 读取图像
        img = cv2.imread(image_path)
        if img is None:
            raise ValueError("无法读取图像，请检查路径是否正确")
        
        h, w = img.shape[:2]
        
        # 计算旋转后的正方形四个顶点坐标
        half_side = side_length / 2
        angle_rad = math.radians(angle_deg)
        
        # 原始正方形四个顶点(未旋转)
        corners = np.array([
            [-half_side, -half_side],
            [half_side, -half_side],
            [half_side, half_side],
            [-half_side, half_side]
        ])
        
        # 旋转矩阵
        rot_mat = np.array([
            [math.cos(angle_rad), -math.sin(angle_rad)],
            [math.sin(angle_rad), math.cos(angle_rad)]
        ])
        
        # 旋转后的顶点
        rotated_corners = np.dot(corners, rot_mat.T)
        
        # 转换为图像坐标
        rotated_corners[:, 0] += center_x
        rotated_corners[:, 1] += center_y
        
        # 计算旋转后的边界框
        x_coords = rotated_corners[:, 0]
        y_coords = rotated_corners[:, 1]
        min_x, max_x = int(np.floor(min(x_coords))), int(np.ceil(max(x_coords)))
        min_y, max_y = int(np.floor(min(y_coords))), int(np.ceil(max(y_coords)))
        
        # 创建足够大的空白画布(白色背景)
        canvas_size = max(max_x - min_x, max_y - min_y)
        canvas = np.ones((canvas_size, canvas_size, 3), dtype=np.uint8) * 255
        
        # 计算在画布中的偏移量
        offset_x = (canvas_size - (max_x - min_x)) // 2
        offset_y = (canvas_size - (max_y - min_y)) // 2
        
        # 从原图中截取区域
        src_rect = np.array([
            [min_x, min_y],
            [max_x, min_y],
            [max_x, max_y],
            [min_x, max_y]
        ], dtype=np.float32)
        
        dst_rect = np.array([
            [offset_x, offset_y],
            [offset_x + (max_x - min_x), offset_y],
            [offset_x + (max_x - min_x), offset_y + (max_y - min_y)],
            [offset_x, offset_y + (max_y - min_y)]
        ], dtype=np.float32)
        
        # 计算透视变换矩阵
        M = cv2.getPerspectiveTransform(src_rect, dst_rect)
        
        # 执行透视变换
        warped = cv2.warpPerspective(img, M, (canvas_size, canvas_size), 
                                    borderMode=cv2.BORDER_CONSTANT, 
                                    borderValue=(255, 255, 255))
        
        # 将变换后的图像叠加到画布上
        mask = (warped != 255).any(axis=2)  # 修改为any而不是all
        canvas[mask] = warped[mask]
        
        # 现在我们需要从画布中提取旋转后的正方形
        # 计算正方形在画布中的中心
        square_center_x = offset_x + (center_x - min_x)
        square_center_y = offset_y + (center_y - min_y)
        
        # 计算旋转后的正方形在画布中的坐标
        square_corners = np.array([
            [square_center_x - half_side, square_center_y - half_side],
            [square_center_x + half_side, square_center_y - half_side],
            [square_center_x + half_side, square_center_y + half_side],
            [square_center_x - half_side, square_center_y + half_side]
        ], dtype=np.float32)
        
        # 旋转正方形到水平位置
        rot_mat_inv = cv2.getRotationMatrix2D((square_center_x, square_center_y), -angle_deg, 1.0)
        rotated_square = cv2.warpAffine(canvas, rot_mat_inv, (canvas_size, canvas_size), 
                                    flags=cv2.INTER_LINEAR, 
                                    borderMode=cv2.BORDER_CONSTANT, 
                                    borderValue=(255, 255, 255))
        
        # 截取水平正方形
        x1 = int(square_center_x - half_side)
        y1 = int(square_center_y - half_side)
        x2 = int(square_center_x + half_side)
        y2 = int(square_center_y + half_side)
        
        result = rotated_square[max(0, y1):min(canvas_size, y2), max(0, x1):min(canvas_size, x2)]
        result = cv2.rotate(result, cv2.ROTATE_90_COUNTERCLOCKWISE)
        _, result = cv2.threshold(result, 127, 255, cv2.THRESH_BINARY)
        
        # 直接保存BGR图像（不需要额外转换）
        cv2.imwrite(output_path, result)


    def run(self):
        scenes = [
            dir 
            for dir in os.listdir(self.scenes_dir) 
            if os.path.isdir(os.path.join(self.scenes_dir, dir)) and "scene_" in dir
        ]
        scenes = sorted(scenes, key=lambda x: int(x.split("_")[1]))
        
        dataset_real_path = os.path.join(self.dataset_save_dir, "real")
        os.makedirs(dataset_real_path, exist_ok=True)
        
        csv_file_path = os.path.join(self.dataset_save_dir, "image_data.csv")
        csv_file = open(csv_file_path, "w", newline="")
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(["Scene", "Image", "X", "Y", "RotationAngle"])
        
        for scene_count, scene in enumerate(scenes):
            raw_image_path = os.path.join(self.scenes_dir, scene, f"{scene}.png")
            safe_area_path = os.path.join(self.scenes_dir, scene, f"{scene}_safe.png")
            safe_image = cv2.imread(safe_area_path, cv2.IMREAD_UNCHANGED)
            green_pixels = np.argwhere(
                (safe_image[:, :, 0] == 0) & 
                (safe_image[:, :, 1] == 255) & 
                (safe_image[:, :, 2] == 0)
            )
            img_count = 0
            for img_count in tqdm(
                range(self.dataset_pre_scene),
                desc=f"Processing: {scene_count * self.dataset_pre_scene + img_count + 1}/{len(scenes) * self.dataset_pre_scene}, Scene: {scene}",
                leave=True,  # 保留进度条
            ):
                random_pixel = random.choice(green_pixels)
                angle = np.random.randint(-180, 180)
                pixel_y, pixel_x = int(random_pixel[0]), int(random_pixel[1])
                real_x = (pixel_x - self.scene_size_pixel[0] / 2) * self.resolution
                real_y = (-pixel_y + self.scene_size_pixel[1] / 2) * self.resolution
                
                real_x_center = real_x + self.scan_width / 2 * np.cos(np.deg2rad(angle))
                real_y_center = real_y + self.scan_width / 2 * np.sin(np.deg2rad(angle))
                
                pixel_x_center = int((real_x_center + self.scene_size[0] / 2) / self.resolution)
                pixel_y_center = int((-real_y_center + self.scene_size[1] / 2) / self.resolution)
                
                output_image_name = str(scene_count * self.dataset_pre_scene + img_count + 1).zfill(5) + ".png"
                output_image = os.path.join(
                    dataset_real_path, 
                    output_image_name
                )
                
                self.extract_rotated_square(
                    raw_image_path, 
                    pixel_x_center, pixel_y_center, 
                    self.roi_imge_width, angle, 
                    output_image
                )
                
                csv_writer.writerow(
                    [scenes[scene_count], output_image_name, real_x, real_y, angle]
                )


if __name__ == "__main__":
    genDataset = GenDataset(
        scenes_dir = "dataset_tools/sources/mjcf_scenes",
        dataset_save_dir = "dataset",
        resolution = 0.02,
        dataset_pre_scene = 10,
    )
    genDataset.run()
    

    
